#!/usr/bin/env python3
"""
RISC-V向量加法演示程序

这个示例展示如何使用RISC-V代码生成器来实现向量加法操作。
演示了以下功能：
1. 数据定义
2. 循环结构
3. 内存访问
4. 寄存器管理
"""

import sys
sys.path.append('..')

from riscv_codegen import RISCVCodeGenerator


def generate_vector_add_basic():
    """生成基本的向量加法代码"""
    print("=== 基本向量加法演示 ===")
    
    # 创建代码生成器
    codegen = RISCVCodeGenerator()
    
    # 定义数据段
    codegen.define_data_word("vector_a", 1, 2, 3, 4, 5)
    codegen.define_data_word("vector_b", 6, 7, 8, 9, 10)
    codegen.define_data_space("vector_result", 20)  # 5个元素 * 4字节
    codegen.define_data_word("vector_length", 5)
    
    # 生成主程序
    codegen.emit_label("_start")
    codegen.emit_comment("向量加法演示程序")
    
    # 使用low-level API实现向量加法
    # 加载数组地址
    base_a = codegen.allocate_register('t0')
    base_b = codegen.allocate_register('t1')
    base_result = codegen.allocate_register('t2')
    counter = codegen.allocate_register('t3')
    length = codegen.allocate_register('t4')
    
    codegen.instructions.la(base_a, "vector_a")
    codegen.instructions.la(base_b, "vector_b")
    codegen.instructions.la(base_result, "vector_result")
    codegen.instructions.lw(length, 0, codegen.register_manager.get_register('gp'))  # 假设长度在全局
    codegen.instructions.li(length, 5)  # 直接设置长度为5
    
    # 使用循环结构
    with codegen.create_loop(counter_reg=counter, limit_reg=length):
        # 计算偏移量 (counter * 4)
        offset = codegen.allocate_register('t5')
        codegen.instructions.slli(offset, counter, 2)
        
        # 加载操作数
        val_a = codegen.allocate_register('a0')
        val_b = codegen.allocate_register('a1')
        temp_addr = codegen.allocate_register('t6')
        
        # 从vector_a加载
        codegen.instructions.add(temp_addr, base_a, offset)
        codegen.instructions.lw(val_a, 0, temp_addr)
        
        # 从vector_b加载
        codegen.instructions.add(temp_addr, base_b, offset)
        codegen.instructions.lw(val_b, 0, temp_addr)
        
        # 执行加法
        codegen.instructions.add(val_a, val_a, val_b)
        
        # 存储结果
        codegen.instructions.add(temp_addr, base_result, offset)
        codegen.instructions.sw(val_a, 0, temp_addr)
        
        # 释放临时寄存器
        codegen.free_register(offset)
        codegen.free_register(val_a)
        codegen.free_register(val_b)
        codegen.free_register(temp_addr)
    
    # 程序结束
    codegen.instructions.li('a0', 0)      # 退出码
    codegen.instructions.li('a7', 93)     # exit系统调用号
    codegen.instructions.ecall()          # 系统调用（需要添加到指令集）
    
    # 释放寄存器
    for reg in [base_a, base_b, base_result, counter, length]:
        codegen.free_register(reg)
    
    return codegen


def generate_vector_add_optimized():
    """生成优化的向量加法代码，使用高级API"""
    print("\n=== 优化向量加法演示 ===")
    
    codegen = RISCVCodeGenerator()
    
    # 定义数据
    codegen.define_data_word("vec_x", *range(1, 11))    # [1,2,3,4,5,6,7,8,9,10]
    codegen.define_data_word("vec_y", *range(11, 21))   # [11,12,13,14,15,16,17,18,19,20]
    codegen.define_data_space("vec_z", 40)              # 10个元素 * 4字节
    
    # 主程序开始
    codegen.emit_label("_start")
    codegen.emit_comment("优化的向量加法: z = x + y")
    
    # 使用高级向量加法API
    with codegen.vector_add("vec_z", "vec_x", "vec_y", 10) as vector_op:
        vector_op.generate_loop()
    
    # 添加系统调用指令到指令集
    codegen.emit("li a0, 0")
    codegen.emit("li a7, 93") 
    codegen.emit("ecall")
    
    return codegen


def generate_conditional_vector_add():
    """生成带条件判断的向量加法"""
    print("\n=== 条件向量加法演示 ===")
    
    codegen = RISCVCodeGenerator()
    
    # 定义数据
    codegen.define_data_word("input_a", -2, -1, 0, 1, 2, 3, -5, 4, -1, 6)
    codegen.define_data_word("input_b", 1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
    codegen.define_data_space("output", 40)
    
    codegen.emit_label("_start")
    codegen.emit_comment("条件向量加法: 只有当a[i] >= 0时，才计算result[i] = a[i] + b[i]")
    
    # 分配寄存器
    base_a = codegen.allocate_register('s0')
    base_b = codegen.allocate_register('s1')
    base_result = codegen.allocate_register('s2')
    counter = codegen.allocate_register('t0')
    length = codegen.allocate_register('t1')
    
    # 初始化
    codegen.instructions.la(base_a, "input_a")
    codegen.instructions.la(base_b, "input_b")
    codegen.instructions.la(base_result, "output")
    codegen.instructions.li(length, 10)
    
    # 循环处理每个元素
    with codegen.create_loop(counter_reg=counter, limit_reg=length):
        # 计算当前偏移量
        offset = codegen.allocate_register('t2')
        codegen.instructions.slli(offset, counter, 2)
        
        # 加载a[i]
        val_a = codegen.allocate_register('t3')
        temp_addr = codegen.allocate_register('t4')
        codegen.instructions.add(temp_addr, base_a, offset)
        codegen.instructions.lw(val_a, 0, temp_addr)
        
        # 检查a[i]是否 >= 0
        if_ctx = codegen.create_if(val_a, 'ge', 'zero')
        with if_ctx:
            # a[i] >= 0, 执行加法
            val_b = codegen.allocate_register('t5')
            codegen.instructions.add(temp_addr, base_b, offset)
            codegen.instructions.lw(val_b, 0, temp_addr)
            
            # 计算结果
            codegen.instructions.add(val_a, val_a, val_b)
            
            # 存储结果
            codegen.instructions.add(temp_addr, base_result, offset)
            codegen.instructions.sw(val_a, 0, temp_addr)
            
            codegen.free_register(val_b)
            
            # 进入else块
            if_ctx.else_block()
            # a[i] < 0, 存储0
            codegen.instructions.add(temp_addr, base_result, offset)
            codegen.instructions.sw('zero', 0, temp_addr)
        
        # 释放临时寄存器
        codegen.free_register(offset)
        codegen.free_register(val_a)
        codegen.free_register(temp_addr)
    
    # 程序结束
    codegen.emit("li a0, 0")
    codegen.emit("li a7, 93")
    codegen.emit("ecall")
    
    return codegen


def generate_function_based_vector_add():
    """生成基于函数的向量加法"""
    print("\n=== 函数式向量加法演示 ===")
    
    codegen = RISCVCodeGenerator()
    
    # 定义数据
    codegen.define_data_word("arr1", *range(1, 6))  # [1,2,3,4,5]
    codegen.define_data_word("arr2", *range(6, 11)) # [6,7,8,9,10]
    codegen.define_data_space("result", 20)         # 5个元素
    
    # 定义向量加法函数
    with codegen.create_function("vector_add_func"):
        # 函数参数：a0=dst, a1=src1, a2=src2, a3=length
        codegen.emit_comment("函数参数: a0=dst, a1=src1, a2=src2, a3=length")
        
        counter = codegen.allocate_register('t0')
        
        with codegen.create_loop(counter_reg=counter, limit_reg='a3'):
            # 计算偏移量
            offset = codegen.allocate_register('t1')
            codegen.instructions.slli(offset, counter, 2)
            
            # 加载操作数
            val1 = codegen.allocate_register('t2')
            val2 = codegen.allocate_register('t3')
            addr = codegen.allocate_register('t4')
            
            # 加载src1[i]
            codegen.instructions.add(addr, 'a1', offset)
            codegen.instructions.lw(val1, 0, addr)
            
            # 加载src2[i]
            codegen.instructions.add(addr, 'a2', offset)
            codegen.instructions.lw(val2, 0, addr)
            
            # 计算和
            codegen.instructions.add(val1, val1, val2)
            
            # 存储到dst[i]
            codegen.instructions.add(addr, 'a0', offset)
            codegen.instructions.sw(val1, 0, addr)
            
            # 释放寄存器
            for reg in [offset, val1, val2, addr]:
                codegen.free_register(reg)
    
    # 主程序
    codegen.emit_label("_start")
    codegen.emit_comment("调用向量加法函数")
    
    # 准备函数参数
    codegen.instructions.la('a0', "result")
    codegen.instructions.la('a1', "arr1")
    codegen.instructions.la('a2', "arr2")
    codegen.instructions.li('a3', 5)
    
    # 调用函数
    codegen.instructions.jal('ra', "vector_add_func")
    
    # 程序结束
    codegen.emit("li a0, 0")
    codegen.emit("li a7, 93")
    codegen.emit("ecall")
    
    return codegen


def main():
    """主函数，演示所有示例"""
    print("RISC-V汇编代码生成器演示")
    print("=" * 50)
    
    # 生成各种示例
    examples = [
        ("basic", generate_vector_add_basic),
        ("optimized", generate_vector_add_optimized),
        ("conditional", generate_conditional_vector_add),
        ("function", generate_function_based_vector_add)
    ]
    
    for name, generator in examples:
        try:
            codegen = generator()
            
            # 保存到文件
            filename = f"vector_add_{name}.s"
            codegen.save_to_file(filename)
            print(f"已保存到文件: {filename}")
            
            # 显示生成的代码（前30行）
            code = codegen.generate_code()
            lines = code.split('\n')
            print(f"\n生成的代码预览 (前30行):")
            print("-" * 40)
            for i, line in enumerate(lines[:30]):
                print(f"{i+1:2d}: {line}")
            if len(lines) > 30:
                print(f"... (还有{len(lines)-30}行)")
            
            print("\n" + "=" * 50)
            
        except Exception as e:
            print(f"生成{name}示例时出错: {e}")
            import traceback
            traceback.print_exc()


if __name__ == "__main__":
    main() 